#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
树形结构Neo4j查询工具
功能：基于树形结构进行自然语言查询
特点：
1. 支持自然语言查询
2. 树形路径查询
3. 关键词匹配
4. 产品范围限定
"""

from neo4j import GraphDatabase
import jieba
import re
import sys
import hashlib
from collections import defaultdict


class TreeQueryTool:
    def __init__(self, uri="bolt://192.168.1.169:7687", user="neo4j", password="Pa@ssw0rd"):
        try:
            self.driver = GraphDatabase.driver(uri, auth=(user, password))
            print(f"✅ 成功连接到Neo4j数据库: {uri}")
        except Exception as e:
            print(f"❌ 连接Neo4j失败: {e}")
            self.driver = None
    
    def close(self):
        if self.driver:
            self.driver.close()
    
    def extract_product_code(self, query):
        """从查询中提取产品编号"""
        # 匹配产品编号模式 (如 TD504, PK511等)
        product_match = re.search(r'([A-Z]{1,3}\d{3,4})', query.upper())
        return product_match.group(1) if product_match else None
    
    def segment_query(self, query):
        """分词并提取关键词"""
        # 移除产品编号
        clean_query = re.sub(r'[A-Z]{1,3}\d{3,4}', '', query, flags=re.IGNORECASE)
        
        # 分词
        words = list(jieba.cut(clean_query))
        
        # 过滤关键词
        keywords = []
        for word in words:
            word = word.strip()
            if len(word) > 1 and word not in ['的', '是', '怎么', '样', '什么', '如何']:
                keywords.append(word)
        
        return keywords
    
    def  search_tree_paths(self, query, limit=None):
        """搜索树形路径"""
        if not self.driver:
            return []

        # 提取产品编号和关键词
        product_code = self.extract_product_code(query)
        keywords = self.segment_query(query)

        print(f"🔍 查询分析:")
        print(f"   产品编号: {product_code or '未指定'}")
        print(f"   关键词: {keywords}")

        results = []

        with self.driver.session() as session:
            # 简化查询策略：直接搜索所有节点
            params = {}

            # 构建查询条件
            where_conditions = []

            if product_code:
                where_conditions.append("n.product_code = $product_code")
                params['product_code'] = product_code

            # 关键词匹配 - 同时搜索内容和关键词
            if keywords:
                keyword_conditions = []
                for i, keyword in enumerate(keywords):
                    content_param = f"content_keyword_{i}"
                    keyword_param = f"keyword_{i}"

                    # 在内容中搜索
                    keyword_conditions.append(f"n.content CONTAINS ${content_param}")
                    params[content_param] = keyword

                    # 在关键词数组中搜索
                    keyword_conditions.append(f"ANY(k IN n.keywords WHERE k CONTAINS ${keyword_param})")
                    params[keyword_param] = keyword

                if keyword_conditions:
                    where_conditions.append(f"({' OR '.join(keyword_conditions)})")

            where_clause = " AND ".join(where_conditions) if where_conditions else "true"

            # 只查询匹配的树节点，不查询叶子节点，不限制数量
            tree_query = f"""
                MATCH (n:TreeNode)
                WHERE {where_clause}
                RETURN DISTINCT n
            """
            print(tree_query)
            try:
                # 只搜索树节点
                tree_result = session.run(tree_query, params)
                matched_tree_nodes = [record['n'] for record in tree_result]

                # 使用集合去重，基于节点ID
                seen_nodes = set()

                # 只返回匹配的树节点，去重
                for node in matched_tree_nodes:
                    node_id = node.get('id')
                    if node_id not in seen_nodes:
                        seen_nodes.add(node_id)
                        results.append({
                            'product_code': node.get('product_code'),
                            'node_type': 'TreeNode',
                            'content': node.get('content'),
                            'filename': node.get('filename', ''),
                            'line_number': node.get('line_number', 0),
                            'keywords': node.get('keywords', []),
                            'node_id': node_id
                        })

            except Exception as e:
                print(f"❌ 查询执行失败: {e}")
                print(f"查询语句: {tree_query}")
                print(f"参数: {params}")

            # 计算相关性并排序
            for result in results:
                result['relevance_score'] = self._calculate_node_relevance(result, keywords)

            results.sort(key=lambda x: x['relevance_score'], reverse=True)

            # 如果去重后只有一个节点，或者只有一个高相关性节点，查询该节点的所有链路
            high_relevance_results = [r for r in results if r['relevance_score'] >= 0.8]

            if len(results) == 1 or len(high_relevance_results) == 1:
                target_node = results[0] if len(results) == 1 else high_relevance_results[0]
                print(f"🔍 找到唯一匹配节点，正在查询其所有链路...")
                chain_results = self._get_node_chains(session, target_node)
                return chain_results

        return results

    def search_by_node_name(self, node_name, product_code=None):
        """根据节点名称查询其一级子节点"""
        if not self.driver:
            return []

        print(f"🔍 查询节点 '{node_name}' 的子节点")
        if product_code:
            print(f"   限定产品: {product_code}")

        results = []

        with self.driver.session() as session:
            # 构建查询条件，找到父节点
            where_conditions = ["parent.content = $node_name"]
            params = {'node_name': node_name}

            if product_code:
                where_conditions.append("parent.product_code = $product_code")
                params['product_code'] = product_code

            where_clause = " AND ".join(where_conditions)

            # 查询父节点的所有一级子节点
            children_query = f"""
                MATCH (parent:TreeNode)
                WHERE {where_clause}
                OPTIONAL MATCH (parent)-[:LEADS_TO]->(child:TreeNode)
                OPTIONAL MATCH (parent)-[:RESULTS_IN]->(leaf:LeafNode)
                RETURN DISTINCT parent, child, leaf
            """

            try:
                children_result = session.run(children_query, params)

                parent_nodes = []
                child_nodes = []
                leaf_nodes = []

                for record in children_result:
                    parent = record['parent']
                    child = record['child']
                    leaf = record['leaf']

                    if parent and parent not in parent_nodes:
                        parent_nodes.append(parent)

                    if child:
                        child_nodes.append(child)

                    if leaf:
                        leaf_nodes.append(leaf)

                if not parent_nodes:
                    print("❌ 没有找到匹配的父节点")
                    return []

                # 去重子节点
                unique_children = []
                seen_contents = set()

                # 添加树节点子节点
                for child in child_nodes:
                    content = child.get('content')
                    if content and content not in seen_contents:
                        seen_contents.add(content)
                        unique_children.append({
                            'product_code': child.get('product_code'),
                            'node_type': 'TreeNode',
                            'content': content,
                            'filename': child.get('filename', ''),
                            'line_number': child.get('line_number', 0),
                            'keywords': child.get('keywords', []),
                            'node_id': child.get('id'),
                            'relevance_score': 1.0
                        })

                # 添加叶子节点子节点
                for leaf in leaf_nodes:
                    content = leaf.get('content')
                    if content and content not in seen_contents:
                        seen_contents.add(content)
                        unique_children.append({
                            'product_code': leaf.get('product_code'),
                            'node_type': 'LeafNode',
                            'content': content,
                            'filename': leaf.get('filename', ''),
                            'line_number': leaf.get('line_number', 0),
                            'keywords': leaf.get('keywords', []),
                            'node_id': leaf.get('id'),
                            'relevance_score': 1.0
                        })

                print(f"✅ 找到 {len(unique_children)} 个子节点")

                # 如果只有一个子节点，返回从父节点开始的全链路
                if len(unique_children) == 1:
                    print(f"🔍 只有一个子节点，正在查询全链路...")
                    parent_node = parent_nodes[0]
                    chain_results = self._get_node_chains(session, {
                        'node_id': parent_node['id'],
                        'product_code': parent_node.get('product_code'),
                        'filename': parent_node.get('filename', ''),
                        'line_number': parent_node.get('line_number', 0)
                    })
                    return chain_results

                # 否则返回所有一级子节点
                return unique_children

            except Exception as e:
                print(f"❌ 子节点查询失败: {e}")
                return []

    def _get_node_chains(self, session, node):
        """获取节点的所有链路"""
        node_id = node['node_id']
        product_code = node['product_code']

        chains = []

        try:
            # 优化1: 分步查询，先找到从根到目标节点的路径
            root_to_target_query = """
                MATCH (p:Product {code: $product_code})-[:HAS_ROOT]->(root:TreeNode)
                MATCH path = (root)-[:LEADS_TO*0..]->(target:TreeNode {id: $node_id})
                RETURN nodes(path) as root_path
                ORDER BY length(path) ASC
                LIMIT 1
            """

            root_result = session.run(root_to_target_query, {
                'product_code': product_code,
                'node_id': node_id
            })

            root_record = root_result.single()
            if not root_record:
                print("❌ 未找到从根到目标节点的路径")
                return [node]

            root_path = root_record['root_path']

            # 优化查询：分步查询提高性能
            # 第一步：获取从根到目标节点的路径
            root_to_target_query = """
                MATCH (p:Product {code: $product_code})-[:HAS_ROOT]->(root:TreeNode)
                MATCH path = (root)-[:LEADS_TO*0..]->(target:TreeNode {id: $node_id})
                RETURN nodes(path) as root_to_target_path
                ORDER BY length(path) ASC
                LIMIT 1
            """

            root_result = session.run(root_to_target_query, {
                'product_code': product_code,
                'node_id': node_id
            })

            root_record = root_result.single()
            if not root_record:
                print("❌ 未找到从根到目标节点的路径")
                return [node]

            root_to_target_path = root_record['root_to_target_path']

            # 第二步：获取从目标节点到所有叶子节点的路径（移除深度限制）
            target_to_leaves_query = """
                MATCH (target:TreeNode {id: $node_id})
                MATCH path = (target)-[:LEADS_TO*0..]->(end)-[:RESULTS_IN]->(leaf:LeafNode)
                WHERE leaf.line_number IS NOT NULL
                RETURN DISTINCT nodes(path) as target_to_leaf_path, leaf, leaf.line_number as line_number
                ORDER BY leaf.line_number
            """

            target_to_leaves_result = session.run(target_to_leaves_query, {'node_id': node_id})

            # 收集所有记录并进行批量处理
            all_records = list(target_to_leaves_result)
            print(f"🔍 找到 {len(all_records)} 条潜在路径")

            # 基于阈值控制处理数量
            max_process = min(len(all_records), 100)  # 最多处理100条
            if len(all_records) > max_process:
                print(f"⚠️  路径数量较多，处理前{max_process}条")
                all_records = all_records[:max_process]

            # 基于完整路径内容的MD5哈希去重，确保每个唯一的决策路径只返回一次
            seen_path_hashes = set()

            for record in all_records:
                target_to_leaf_path = record['target_to_leaf_path']
                leaf = record['leaf']
                line_number = record['line_number']

                # 构建完整路径内容用于哈希
                full_path_content = []

                # 添加从根到目标节点的路径内容
                for path_node in root_to_target_path:
                    full_path_content.append(path_node['content'])

                # 添加从目标到叶子的路径内容（排除重复的目标节点）
                if len(target_to_leaf_path) > 1:
                    for path_node in target_to_leaf_path[1:]:
                        full_path_content.append(path_node['content'])

                # 添加叶子节点内容
                if leaf['content'] not in full_path_content:
                    full_path_content.append(leaf['content'])

                # 生成路径内容的MD5哈希
                path_content_str = " -> ".join(full_path_content)
                path_hash = hashlib.md5(path_content_str.encode('utf-8')).hexdigest()

                # 基于MD5哈希去重
                if path_hash in seen_path_hashes:
                    continue
                seen_path_hashes.add(path_hash)

                # 构建完整链路（组合根到目标 + 目标到叶子）
                complete_chain = []

                # 添加从根到目标节点的路径
                for path_node in root_to_target_path:
                    complete_chain.append({
                        'content': path_node['content'],
                        'node_type': 'TreeNode',
                        'is_matched': path_node['id'] == node_id
                    })

                # 添加从目标到叶子的路径（排除重复的目标节点）
                if len(target_to_leaf_path) > 1:
                    for path_node in target_to_leaf_path[1:]:
                        complete_chain.append({
                            'content': path_node['content'],
                            'node_type': 'TreeNode',
                            'is_matched': False
                        })

                # 只有当叶子节点内容与最后一个树节点不同时才添加叶子节点
                if not complete_chain or leaf['content'] != complete_chain[-1]['content']:
                    complete_chain.append({
                        'content': leaf['content'],
                        'node_type': 'LeafNode',
                        'is_matched': False
                    })

                chains.append({
                    'product_code': product_code,
                    'chain': complete_chain,
                    'matched_node': node,
                    'filename': node['filename'],
                    'line_number': node['line_number']
                })

            print(f"✅ 找到 {len(chains)} 条链路")
            return chains

        except Exception as e:
            print(f"❌ 查询链路失败: {e}")
            return [node]  # 返回原节点

    def _calculate_node_relevance(self, node, keywords):
        """计算节点与查询的相关性"""
        if not keywords:
            return 1.0

        score = 0.0
        content = node.get('content', '')
        node_keywords = node.get('keywords', [])

        for keyword in keywords:
            # 在内容中匹配
            if keyword in content:
                score += 1.0
            # 在关键词中匹配
            elif any(keyword in k for k in node_keywords):
                score += 0.8

        return score / len(keywords) if keywords else 0.0
    
    def display_results(self, results):
        """显示查询结果"""
        if not results:
            print("❌ 没有找到匹配的结果")
            return

        print(f"\n🎯 找到 {len(results)} 个相关结果:")
        print("=" * 80)

        for i, result in enumerate(results, 1):
            # 检查是否是链路结果
            if 'chain' in result:
                # 显示链路
                product_code = result['product_code']
                chain = result['chain']
                filename = result['filename']
                line_number = result['line_number']

                print(f"\n📋 链路 {i} (产品: {product_code})")
                print(f"📁 来源: {filename}:{line_number}")
                print("🔗 决策链路:")

                for j, step in enumerate(chain):
                    content = step['content']
                    node_type = step['node_type']
                    is_matched = step.get('is_matched', False)

                    # 节点类型图标
                    if node_type == 'LeafNode':
                        icon = "🎯"
                        prefix = "   └─"
                    else:
                        icon = "🔍" if is_matched else "📍"
                        if j == len(chain) - 1:  # 最后一个
                            prefix = "   └─"
                        else:
                            prefix = "   ├─"

                    print(f"{prefix} {icon} {content}")

                print("-" * 60)
            else:
                # 显示普通节点
                product_code = result['product_code']
                content = result['content']
                relevance = result['relevance_score']
                filename = result['filename']
                line_number = result['line_number']

                print(f"\n📋 结果 {i} (🌳 决策节点)")
                print(f"🏷️  产品: {product_code}")
                print(f"📄 内容: {content}")
                print(f"📁 来源: {filename}:{line_number}")
                print(f"⭐ 相关性: {relevance:.2f}")
                print("-" * 60)
    
    def interactive_search(self):
        """交互式搜索"""
        print("\n🔍 树形结构交互式查询工具")
        print("💡 支持两种查询模式:")
        print("   1. 自然语言查询: 'TD504投入量', 'PK511扩增循环数', '温度设置'")
        print("   2. 节点名称查询: 'node:节点名称' 或 'node:节点名称:产品编号'")
        print("💡 输入 'quit' 或 'exit' 退出")
        print("=" * 60)

        while True:
            try:
                query = input("\n🔍 请输入查询: ").strip()

                if query.lower() in ['quit', 'exit', 'q']:
                    print("👋 再见！")
                    break

                if not query:
                    continue

                # 检查是否是节点名称查询
                if query.startswith('node:'):
                    parts = query[5:].split(':')
                    node_name = parts[0].strip()
                    product_code = parts[1].strip() if len(parts) > 1 else None

                    print(f"\n🔄 正在查询节点: {node_name}")
                    if product_code:
                        print(f"   限定产品: {product_code}")

                    results = self.search_by_node_name(node_name, product_code)
                    self.display_results(results)
                else:
                    # 自然语言查询
                    print(f"\n🔄 正在搜索: {query}")
                    results = self.search_tree_paths(query)
                    self.display_results(results)

            except KeyboardInterrupt:
                print("\n👋 再见！")
                break
            except Exception as e:
                print(f"❌ 查询出错: {e}")


def main():
    """主函数"""
    print("🌳 树形结构Neo4j查询工具")
    print("=" * 50)
    
    # 创建查询工具
    query_tool = TreeQueryTool()
    
    if not query_tool.driver:
        print("❌ 无法连接到Neo4j数据库")
        return
    
    try:
        # 检查是否有命令行参数
        if len(sys.argv) > 1:
            # 直接查询模式
            query = ' '.join(sys.argv[1:])
            print(f"🔍 查询: {query}")
            results = query_tool.search_tree_paths(query)
            query_tool.display_results(results)
        else:
            # 交互式模式
            query_tool.interactive_search()
    
    finally:
        query_tool.close()


if __name__ == "__main__":
    main()
